package com.wayhome.srpingbootmybatis.metrics.async;

import com.wayhome.srpingbootmybatis.async.TaskListener;
import com.wayhome.srpingbootmybatis.metrics.PrefixFilteredCollectorRegistry;
import io.prometheus.client.Counter;
import io.prometheus.client.Gauge;
import io.prometheus.client.Histogram;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AsyncTaskMetrics {


    @Bean
    public TaskMetricsListener taskMetricsListener() {
        return new TaskMetricsListener();
    }


    public static class TaskMetricsListener implements TaskListener {

        private static Histogram histogram = Histogram.build("async_task", "任务执行时间分布").labelNames("type").register(PrefixFilteredCollectorRegistry.getInstance());
        private static Gauge gauge     = Gauge.build("async_task_running_num","当前正在进行的任务数量").labelNames("type").register(PrefixFilteredCollectorRegistry.getInstance());
        private static Counter errorCounter     = Counter.build("async_task_error_num","运行报错的任务数量累计").labelNames("type").register(PrefixFilteredCollectorRegistry.getInstance());

        private static ThreadLocal<Histogram.Timer> timerThreadLocal = new ThreadLocal<>();

        @Override
        public void beforeExecute(TaskListener.TaskExecutionContext taskContext) {
            gauge.labels(taskContext.getType()).inc();
            Histogram.Timer timer = histogram.labels(taskContext.getType()).startTimer();
            timerThreadLocal.set(timer);
        }
        @Override
        public void afterExecute(TaskExecutionContext taskContext, Throwable ex) {
            gauge.labels(taskContext.getType()).dec();
            if (ex != null) {
                errorCounter.labels(taskContext.getType()).inc();
            }
            Histogram.Timer timer = timerThreadLocal.get();
            if (timer != null) {
                timer.observeDuration();
            }
            timerThreadLocal.remove();
        }

        @Override
        public void afterSetRunning(TaskListener.TaskExecutionContext taskContext) {
        }

    }


}
